#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <string.h>

typedef struct {
    double *x;
    double *y;
    int n;
} Dataset;

Dataset read_data(const char *filename) {
    FILE *fp = fopen(filename, "r");
    if (!fp) {
        perror("Error opening file");
        exit(1);
    }

    Dataset data;
    data.n = 0;
    data.x = NULL;
    data.y = NULL;
    double x_val, y_val;
    while (fscanf(fp, "%lf %lf", &x_val, &y_val) == 2) {
        data.x = realloc(data.x, (data.n + 1) * sizeof(double));
        data.y = realloc(data.y, (data.n + 1) * sizeof(double));
        data.x[data.n] = x_val;
        data.y[data.n] = y_val;
        data.n++;
    }

    fclose(fp);
    return data;
}

void gauss_elimination(double **A, double *B, int n, double *theta) {
    for (int i = 0; i < n; i++) {
        int max_row = i;
        for (int k = i + 1; k < n; k++) {
            if (fabs(A[k][i]) > fabs(A[max_row][i])) {
                max_row = k;
            }
        }

        if (max_row != i) {
            double *tmp_row = A[i];
            A[i] = A[max_row];
            A[max_row] = tmp_row;

            double tmp = B[i];
            B[i] = B[max_row];
            B[max_row] = tmp;
        }

        if (fabs(A[i][i]) < 1e-10) {
            fprintf(stderr, "Matrix is singular or nearly singular\n");
            exit(1);
        }

        for (int k = i + 1; k < n; k++) {
            double factor = A[k][i] / A[i][i];
            B[k] -= factor * B[i];
            for (int j = i; j < n; j++) {
                A[k][j] -= factor * A[i][j];
            }
        }
    }

    for (int i = n - 1; i >= 0; i--) {
        theta[i] = B[i];
        for (int j = i + 1; j < n; j++) {
            theta[i] -= A[i][j] * theta[j];
        }
        theta[i] /= A[i][i];
    }
}

void least_squares(Dataset data, int m, double *theta) {
    int max_power = 2 * m;
    double *sum_x_powers = calloc(max_power + 1, sizeof(double));
    double *sum_x_powers_y = calloc(m + 1, sizeof(double));

    for (int i = 0; i < data.n; i++) {
        double xi = data.x[i];
        double yi = data.y[i];
        double pow_xi = 1.0;

        for (int k = 0; k <= max_power; k++) {
            sum_x_powers[k] += pow_xi;
            pow_xi *= xi;
        }

        pow_xi = 1.0;
        for (int k = 0; k <= m; k++) {
            sum_x_powers_y[k] += pow_xi * yi;
            pow_xi *= xi;
        }
    }

    int size = m + 1;
    double **A = malloc(size * sizeof(double *));
    for (int i = 0; i < size; i++) {
        A[i] = malloc(size * sizeof(double));
    }
    double *B = malloc(size * sizeof(double));

    for (int i = 0; i < size; i++) {
        for (int j = 0; j < size; j++) {
            A[i][j] = sum_x_powers[i + j]; 
        }
        B[i] = sum_x_powers_y[i];      
    }

    gauss_elimination(A, B, size, theta);

    free(sum_x_powers);
    free(sum_x_powers_y);
    for (int i = 0; i < size; i++) free(A[i]);
    free(A);
    free(B);
}

void gradient_descent(Dataset data, int m, double *theta, double lr, int iterations) {
    int size = m + 1;
    for (int i = 0; i < size; i++) theta[i] = 0.0;

    for (int iter = 0; iter < iterations; iter++) {
        double *grad = calloc(size, sizeof(double));

        for (int i = 0; i < data.n; i++) {
            double xi = data.x[i];
            double yi = data.y[i];
            double y_pred = 0.0;

            for (int j = 0; j < size; j++) {
                y_pred += theta[j] * pow(xi, j); 
            }

            double error = y_pred - yi;
            for (int j = 0; j < size; j++) {
                grad[j] += error * pow(xi, j);
            }
        }

        for (int j = 0; j < size; j++) {
            theta[j] -= lr * grad[j] / data.n;
        }

        free(grad);
    }
}

double calculate_rss(Dataset data, int m, double *theta) {
    double rss = 0.0;
    for (int i = 0; i < data.n; i++) {
        double y_pred = 0.0;
        for (int j = 0; j <= m; j++) {
            y_pred += theta[j] * pow(data.x[i], j); 
        }
        rss += pow(y_pred - data.y[i], 2);
    }
    return rss;
}

void plot_data(Dataset data, int m, double *theta) {
    FILE *data_fp = fopen("data_points.dat", "w");
    for (int i = 0; i < data.n; i++) {
        fprintf(data_fp, "%f %f\n", data.x[i], data.y[i]);
    }
    fclose(data_fp);

    FILE *curve_fp = fopen("curve.dat", "w");
    for (double x = -100; x <= 100; x += 0.5) {
        double y = 0.0;
        for (int j = 0; j <= m; j++) {
            y += theta[j] * pow(x, j); 
        }
        fprintf(curve_fp, "%f %f\n", x, y);
    }
    fclose(curve_fp);

    FILE *gnuplot = popen("gnuplot -persistent", "w");
    fprintf(gnuplot, "plot 'data_points.dat' with points title 'Data', 'curve.dat' with lines title 'Fit'\n");
    fflush(gnuplot);
    pclose(gnuplot);
}

int main(int argc, char *argv[]) {
    if (argc != 4) {
        fprintf(stderr, "Usage: %s data.dat [ls|gd] t\n", argv[0]);
        return 1;
    }

    Dataset data = read_data(argv[1]);
    char *method = argv[2];
    int t = atoi(argv[3]);
    int m = t;

    int size = m + 1;
    double *theta = malloc(size * sizeof(double));

    if (strcmp(method, "ls") == 0) {
        least_squares(data, m, theta);
    } else if (strcmp(method, "gd") == 0) {
        gradient_descent(data, m, theta, 0.0001, 100000);
    } else {
        fprintf(stderr, "Invalid method: %s\n", method);
        return 1;
    }

    printf("Parameters: ");
    for (int i = 0; i < size; i++) {
        printf("%f ", theta[i]);
    }
    printf("\nRSS: %f\n", calculate_rss(data, m, theta));

    plot_data(data, m, theta);

    free(data.x);
    free(data.y);
    free(theta);
    return 0;
}